package com.wilmer.ai.listener;

import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.unfbx.chatgpt.entity.chat.BaseMessage;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.Message;
import com.wilmer.TecherLessonPlan.domain.vo.ChatFinishVo;
import com.wilmer.ai.config.LocalCache;
import com.wilmer.aiMessage.domain.AiMessage;
import com.wilmer.aiMessage.mapper.AiMessageMapper;
import com.wilmer.aiUserMsg.domain.YwUserMsg;
import com.wilmer.aiUserMsg.mapper.YwUserMsgMapper;
import com.wilmer.aiUserMsg.service.IYwUserMsgService;
import com.wilmer.common.core.domain.model.LoginUser;
import com.wilmer.common.utils.SecurityUtils;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Objects;

/**
 * @author hans
 * @date 2024-05-15 23:47:46
 * @description 流式对话Listener
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class userSseStreamFastGptListener extends EventSourceListener {


    private LoginUser loginUser;

    private long tokens;

    private SseEmitter sseEmitter;

    private String message = "";

    private YwUserMsgMapper msgMapper;

    private String chatId;

    private List<Message> messages;


    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        log.info("OpenAI建立sse连接...");
    }

    @SneakyThrows
    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {

        tokens += 1;
        log.info("OpenAI返回数据：{}", data);
        if ("[DONE]".equals(data)) {
            log.info("OpenAI返回数据结束了");
            System.out.println(message);
            //获取当前Ai回复的记录
            Message currentMessage = Message.builder().content(message).role(Message.Role.ASSISTANT).build();
            //将聊天记录写入数据库
            YwUserMsg userMsg = new YwUserMsg();
            userMsg.setUserId(loginUser.getUserId());
            userMsg.setDeptId(loginUser.getDeptId());
            userMsg.setMsgContext(JSONUtil.toJsonStr(currentMessage));
            userMsg.setChatId(chatId);
            msgMapper.insertYwUserMsg(userMsg);
            //聊天记录放入缓存，后期改为数据库
            messages.add(currentMessage);
            LocalCache.CACHE.put("msg" +
                    chatId, JSONUtil.toJsonStr(messages), LocalCache.TIMEOUT);
            //关闭连接
            sseEmitter.complete();
            return;
        }
        ObjectMapper mapper = new ObjectMapper();
        ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class);
        boolean finish = Objects.equals(completionResponse.getChoices().get(0).getFinishReason(), "stop");
        System.out.println(finish);
        String content = completionResponse.getChoices().get(0).getDelta().getContent();
        if (content != null) {
            message += content;
        }
        try {
            if (finish) {
                sseEmitter.send(new ChatFinishVo());
            } else {
                sseEmitter.send(SseEmitter.event().data(completionResponse.getChoices().get(0).getDelta()));
            }
        } catch (Exception e) {
            log.error("sse信息推送失败！");
            eventSource.cancel();
        }
    }


    @Override
    public void onClosed(EventSource eventSource) {
        log.info("流式输出返回值总共{}tokens", tokens() - 2);
        log.info("OpenAI关闭sse连接...");
    }

    @SneakyThrows
    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        if (Objects.isNull(response)) {
            return;
        }
        ResponseBody body = response.body();
        if (Objects.nonNull(body)) {
            log.error("fastgpt sse连接异常data：{}，异常：{}", body.string(), t);
        } else {
            log.error("fastgpt  sse连接异常data：{}，异常：{}", response, t);
        }
        eventSource.cancel();
    }

    /**
     * tokens
     */
    public long tokens() {
        return tokens;
    }


}
